# -*- coding: utf-8 -*-

import os
import sys
import subprocess
import time
import stat
from subprocess import check_call, call

UNIT = 1024
MOUNT_POINT = "mnt"
MOUNT_POINT_efi = "%s/boot/efi" % MOUNT_POINT

def remove(path):
    try:
        os.remove(path)
    except OSError as err:
        if err.errno == 2:
            print("%s does not exist" % path)
            return
        else:
            raise
    print("%s removed" % path)


def prepare():
    if os.geteuid() != 0:
        print("root privilege required")
        sys.exit(1)
    if not os.path.exists("/dev/nbd0"):
        check_call(["modprobe", "nbd"])


def is_block(pathname):
    if os.path.exists(pathname):
        mode = os.stat(pathname).st_mode
        return stat.S_ISBLK(mode)
    else:
        return False


def test_block(devpath):
    TESTDIR = "testmnt"
    TESTFILE = TESTDIR + "/file1"
    print("mounted " + devpath + " on " + TESTDIR + " for testing")
    os.mkdir(TESTDIR)
    rc = call(["mount", devpath, TESTDIR])
    if rc == 0:
        with open(TESTFILE, "w") as f:
            f.write("ABC")
        os.unlink(TESTFILE)
        call(["umount", TESTDIR])
        os.rmdir(TESTDIR)
        print("mount test succeeded")
    else:
        print("mount test failed")


def detach_nbd(nbdpath):
    if os.path.isdir(MOUNT_POINT):
        for dirname in [os.path.join(MOUNT_POINT, "proc"),
                        os.path.join(MOUNT_POINT, "sys"),
                        os.path.join(MOUNT_POINT, "dev"), MOUNT_POINT_efi, MOUNT_POINT]:
            # subprocess.call("lsof | grep mnt\/dev | awk '{print $2}' | uniq | xargs kill", shell = True)
            call(["umount", dirname])
        # os.rmdir(MOUNT_POINT)
    call(["kpartx", "-s", "-d", nbdpath])
    call(["qemu-nbd", "-d", nbdpath])


def mount_proc_dev_sys():
    check_call(["mount", "-t", "proc", "proc", os.path.join(MOUNT_POINT,
                                                           "proc")])
    check_call(["mount", "-t", "sysfs", "sysfs", os.path.join(MOUNT_POINT,
                                                             "sys")])
    check_call(["mount", "--bind", "/dev", os.path.join(MOUNT_POINT, "dev")])


def umount_proc_dev_sys():
    for dirname in [os.path.join(MOUNT_POINT, "proc"),
                    os.path.join(MOUNT_POINT, "sys"),
                    os.path.join(MOUNT_POINT, "dev")]:
        # subprocess.call("lsof | grep mnt\/dev | awk '{print $2}' | uniq | xargs kill", shell = True)
        call(["umount", dirname])


# return devpath. You must run detach_nbd after finishing your job.
def mount_qcow2(filename, arch):
    prepare()
    devpath = ""
    for i in range(1, 16):
        nbdfile = "nbd%d" % i
        devpath_tmp = "/dev/" + nbdfile
        rc = call(["qemu-nbd", "-c", devpath_tmp, filename])
        if rc == 0:
            devpath = devpath_tmp
            break
    if not devpath:
        print("Fail to attach %s to any nbd device" % filename)
        sys.exit(4)

    check_call(["kpartx", "-s", "-a", devpath])

    if arch == 'x86_64':
        partpath = ["/dev/mapper/" + nbdfile + 'p1']
        if not is_block(partpath[0]):
            print("Fail to find " + partpath[0])
            detach_nbd(devpath)
            sys.exit(3)

        try:
            os.mkdir(MOUNT_POINT)
        except OSError as err:
            print("Fail to mkdir %s: %s" % (MOUNT_POINT, str(err)))
            detach_nbd(devpath)
            sys.exit(3)

        if call(["mount", partpath[0], MOUNT_POINT]) != 0:
            print("Fail to mount %s" % partpath[0])
            detach_nbd(devpath)
            sys.exit(3)
        print("mounted " + partpath[0] + " on " + MOUNT_POINT)
    elif arch == 'aarch64':
        partpath = ["/dev/mapper/" + nbdfile + "p1", "/dev/mapper/" + nbdfile + "p2"]
        if not is_block(partpath[0]):
            print("Fail to find " + partpath[0])
            detach_nbd(devpath)
            sys.exit(3)
        if not is_block(partpath[1]):
            print("Fail to find " + partpath[1])
            detach_nbd(devpath)
            sys.exit(3)

        try:
            os.mkdir(MOUNT_POINT)
        except OSError as err:
            print("Fail to mkdir %s: %s" % (MOUNT_POINT, str(err)))
            detach_nbd(devpath)
            sys.exit(3)

        if call(["mount", partpath[1], MOUNT_POINT]) != 0:
            print("Fail to mount %s" % partpath[1])
            detach_nbd(devpath)
            sys.exit(3)

        try:
            os.makedirs(MOUNT_POINT_efi)
        except OSError as err:
            print("Fail to mkdir %s: %s" % (MOUNT_POINT_efi, str(err)))
            detach_nbd(devpath)
            sys.exit(3)

        if call(["mount", partpath[0], MOUNT_POINT_efi]) != 0:
            print("Fail to mount %s" % partpath[0])
            detach_nbd(devpath)
            sys.exit(3)

        print("mounted " + partpath[1] + " on " + MOUNT_POINT)
        print("mounted " + partpath[0] + " on " + MOUNT_POINT_efi)
    elif arch == 'loongarch64':
        partpath = ["/dev/mapper/" + nbdfile + "p1", "/dev/mapper/" + nbdfile + "p2"]
        if not is_block(partpath[0]):
            print("Fail to find " + partpath[0])
            detach_nbd(devpath)
            sys.exit(3)
        if not is_block(partpath[1]):
            print("Fail to find " + partpath[1])
            detach_nbd(devpath)
            sys.exit(3)

        try:
            os.mkdir(MOUNT_POINT)
        except OSError as err:
            print("Fail to mkdir %s: %s" % (MOUNT_POINT, str(err)))
            detach_nbd(devpath)
            sys.exit(3)

        if call(["mount", partpath[1], MOUNT_POINT]) != 0:
            print("Fail to mount %s" % partpath[1])
            detach_nbd(devpath)
            sys.exit(3)

        try:
            os.makedirs(MOUNT_POINT_efi)
        except OSError as err:
            print("Fail to mkdir %s: %s" % (MOUNT_POINT_efi, str(err)))
            detach_nbd(devpath)
            sys.exit(3)

        if call(["mount", partpath[0], MOUNT_POINT_efi]) != 0:
            print("Fail to mount %s" % partpath[0])
            detach_nbd(devpath)
            sys.exit(3)

        print("mounted " + partpath[1] + " on " + MOUNT_POINT)
        print("mounted " + partpath[0] + " on " + MOUNT_POINT_efi)

    return nbdfile


def make_ext4_qcow2(size_in_G, arch, for_data=False):
    prepare()
    size_in_byte = size_in_G * (UNIT**3)
    print("size in byte is %d" % size_in_byte)
    filename = "EXT4_%sG.qcow2" % size_in_G
    if os.path.exists(filename):
        print("%s exists, please remove it and try again" % filename)
        sys.exit(3)
    check_call(["qemu-img", "create", "-f", "qcow2", filename, str(size_in_byte)])
    devpath = ""
    for i in range(0, 16):
        nbdfile = "nbd%d" % i
        devpath_tmp = "/dev/" + nbdfile
        rc = call(["qemu-nbd", "-c", devpath_tmp, filename])
        if rc == 0:
            devpath = devpath_tmp
            break
    if not devpath:
        print("Fail to attach %s to any nbd device" % filename)
        sys.exit(4)
    try:
        if arch == 'x86_64':
            p = subprocess.Popen(["fdisk", devpath], stdin=subprocess.PIPE)
            #p.communicate("n\np\n1\n\n+10G\nn\np\n2\n\n\nw\n")
            p.communicate("n\np\n1\n\n\nw\n".encode())
            p.wait()
            time.sleep(5)
            check_call(["kpartx", "-s", "-a", devpath])

            partpath = ["/dev/mapper/" + nbdfile + 'p1']
            if not is_block(partpath[0]):
                print("Fail to find " + partpath[0])
                detach_nbd(devpath)
                sys.exit(3)
            check_call(["mkfs.xfs", partpath[0]])
        elif arch == 'aarch64':
            p = subprocess.Popen(["gdisk", devpath], stdin=subprocess.PIPE)
            p.communicate("n\n1\n\n+512M\nef00\nn\n2\n\n\n8300\nw\nY\n".encode())
            p.wait()
            time.sleep(5)
            check_call(["kpartx", "-s", "-a", devpath])

            partpath = ["/dev/mapper/" + nbdfile + "p1", "/dev/mapper/" + nbdfile + "p2"]
            # partpath = "/dev/mapper/" + nbdfile + 'p1'
            # partpath1 = "/dev/mapper/" + nbdfile + 'p2'
            if not is_block(partpath[0]):
                print("Fail to find " + partpath[0])
                detach_nbd(devpath)
                sys.exit(3)
            if not is_block(partpath[1]):
                print("Fail to find " + partpath[1])
                detach_nbd(devpath)
                sys.exit(3)

            check_call(["mkfs.fat", "-F", "32", partpath[0]])
            check_call(["mkfs.xfs", partpath[1]])
        elif arch == 'loongarch64':
            p = subprocess.Popen(["gdisk", devpath], stdin=subprocess.PIPE)
            p.communicate("n\n1\n\n+512M\nef00\nn\n2\n\n\n8300\nw\nY\n".encode())
            p.wait()
            time.sleep(5)
            check_call(["kpartx", "-s", "-a", devpath])

            partpath = ["/dev/mapper/" + nbdfile + "p1", "/dev/mapper/" + nbdfile + "p2"]
            # partpath = "/dev/mapper/" + nbdfile + 'p1'
            # partpath1 = "/dev/mapper/" + nbdfile + 'p2'
            if not is_block(partpath[0]):
                print("Fail to find " + partpath[0])
                detach_nbd(devpath)
                sys.exit(3)
            if not is_block(partpath[1]):
                print("Fail to find " + partpath[1])
                detach_nbd(devpath)
                sys.exit(3)

            check_call(["mkfs.fat", "-F", "32", partpath[0]])
            check_call(["mkfs.xfs", partpath[1]])
    except subprocess.CalledProcessError as err:
        print("Subprocess error: " + str(err))
        detach_nbd(devpath)
        sys.exit(5)
    try:
        test_block(partpath[0])
        if arch == 'aarch64':
            test_block(partpath[1])
        elif arch == 'loongarch64':
            test_block(partpath[1])
    finally:
        detach_nbd(devpath)
    return filename


def make_ext3_qcow2(size_in_G, for_data=False):
    prepare()
    size_in_byte = size_in_G * (UNIT**3)
    print("size in byte is %d" % size_in_byte)
    filename = "EXT3_%sG.qcow2" % size_in_G
    if os.path.exists(filename):
        print("%s exists, please remove it and try again" % filename)
        sys.exit(3)
    check_call(["qemu-img", "create", "-f", "qcow2", filename, str(size_in_byte)])
    devpath = ""
    for i in range(0, 16):
        nbdfile = "nbd%d" % i
        devpath_tmp = "/dev/" + nbdfile
        rc = call(["qemu-nbd", "-c", devpath_tmp, filename])
        if rc == 0:
            devpath = devpath_tmp
            break
    if not devpath:
        print("Fail to attach %s to any nbd device" % filename)
        sys.exit(4)
    try:
        p = subprocess.Popen(["fdisk", devpath], stdin=subprocess.PIPE)
        p.communicate("n\n\n1\n\n\na\n\nw\n".encode())
        p.wait()
        time.sleep(5)
        check_call(["kpartx", "-s", "-a", devpath])

        partpath = "/dev/mapper/" + nbdfile + 'p1'

        if not is_block(partpath):
            print("Fail to find " + partpath)
            detach_nbd(devpath)
            sys.exit(3)

        if for_data:
            check_call(["mkfs.ext3", "-N", "6000000", partpath])
        else:
            check_call(["mkfs.ext3", partpath])
    except subprocess.CalledProcessError as err:
        print("Subprocess error: " + str(err))
        detach_nbd(devpath)
        sys.exit(5)
    try:
        test_block(partpath)
    finally:
        detach_nbd(devpath)
    return filename


def grub_install(dirname, devpath):
    devfile = os.path.basename(devpath)
    partpath = os.path.join("/dev/mapper", devfile + "p1")
    fake_link = os.path.join("/dev/", devfile + "p")
    fake_link_part = os.path.join("/dev/", devfile + "p1")
    remove(fake_link)
    os.symlink(devfile, fake_link)
    if not os.path.exists(fake_link_part):
        os.symlink(partpath, fake_link_part)
    print("Update device.map")
    with open(os.path.join(dirname, "boot/grub/device.map"), "w") as f:
        f.write("(hd0)     %s\n" % fake_link )
    grub_setup = subprocess.Popen(["grub", "--batch"], stdin=subprocess.PIPE)
    grub_setup.communicate("""device (hd0) %s
root (hd0,0)
setup (hd0)
quit
""" % fake_link)
    grub_setup.wait()
    os.remove(fake_link)
    if os.path.islink(fake_link_part):
        os.remove(fake_link_part)
    with open(os.path.join(dirname, "boot/grub/device.map"), "w") as f:
        f.write("(hd0)     /dev/vda\n")
